import os
import tensorflow as tf

if __name__ == '__main__':
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    # 创建一个张量
    dataset = tf.reshape(tf.range(15), shape=(5, 3))
    print(dataset)

    # 从最高维度进行切片。结果为 5 个 Dataset
    result_dataset = tf.data.Dataset.from_tensor_slices(dataset)

    # 将 dataset 按照 batch_size 进行合并
    batched_result_dataset = result_dataset.batch(batch_size=3)

    print(batched_result_dataset)

    # for item in batched_result_dataset:
    #     print(item)

    for index, item in enumerate(batched_result_dataset):
        print(f"index: {index} \n item: {item}")

    # 将剩余 不够 batch_size 个数的一批 删除
    batched_result_dataset1 = result_dataset.batch(batch_size=3, drop_remainder=True)

    for index, item in enumerate(batched_result_dataset1):
        print(f"index: {index} \n item: {item}")

    dataset = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]
    prepared_dataset = tf.data.Dataset.from_generator(lambda: iter(dataset), tf.int32)
    # print(dataset)

    # print(list(prepared_dataset))
    print(list(prepared_dataset.as_numpy_iterator()))

    # 按照指定的形状进行补齐，默认值 0
    padded_batch_result_dataset1 = prepared_dataset.padded_batch(2, padded_shapes=4)

    for index, item in enumerate(padded_batch_result_dataset1):
        print(f"index: {index} \n item: {item}")

    # 按照指定的形状进行补齐，指定使用值
    padded_batch_result_dataset2 = prepared_dataset.padded_batch(2, padded_shapes=4, padding_values=9)

    for index, item in enumerate(padded_batch_result_dataset2):
        print(f"index: {index} \n item: {item}")
